import os
import cv2
import sys
import numpy as np
import random
import torch
import torchvision.transforms.functional as F
import matplotlib.pyplot as plt
import time
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
from torch.autograd import Variable
import pyautogui as pg

plt.rcParams['interactive'] = "true"
plt.rcParams["savefig.bbox"] = 'tight'

def getContours(imgA, imgB):
    Borderlist = []  # 不同的轮廓图像及坐标
    x, y, w, h, xx, yy, ss = 0, 0, 10, 10, 20, 20, 10  # 因为图像大小不能为0
    imgGet = np.array([[], []])  # 不能为空
    contours, hierarchy = cv2.findContours(imgA, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)  # 检索外部轮廓
    for cnt in contours:  
        area = cv2.contourArea(cnt)
        if area > 10 and area < 450:  # 面积大于800像素为封闭图形
            cv2.drawContours(imgB, cnt, -1, (255, 0, 0), 3)  
            peri = cv2.arcLength(cnt, True)  # 计算周长
            approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)  # 计算有多少个拐角
            x, y, w, h = cv2.boundingRect(approx)  # 得到外接矩形的大小
            # print("w:%d, h:%d, w/h:%.02f" % (w, h, float(w) / h))
            a = (w + h) // 2
            dd = abs((w - h) // 2)  # 边框的差值
            imgGet = imgA[y:y + h, x:x + w]
            if w <= h:  # 得到一个正方形框，边界往外扩充20像素,黑色边框
                imgGet = cv2.copyMakeBorder(imgGet, 5, 5, 5 + dd, 5 + dd, cv2.BORDER_CONSTANT, value=[0, 0, 0])
                xx = x - dd - 10
                yy = y - 10
                ss = h + 20
                cv2.rectangle(imgB, (x - dd - 5, y - 5), (x + a + 5, y + h + 5), (0, 255, 0), 2)  # 看看框选的效果，在imgCopy中
                print(a + dd, h)
            else:  # 边界往外扩充20像素值
                imgGet = cv2.copyMakeBorder(imgGet, 5 + dd, 5 + dd, 5, 5, cv2.BORDER_CONSTANT, value=[0, 0, 0])
                xx = x - 10
                yy = y - dd - 10
                ss = w + 20
                cv2.rectangle(imgB, (x - 5, y - dd - 5), (x + w + 5, y + a + 5), (0, 255, 0), 2)
                print(a + dd, w)
            # imgGet = cv2.resize(imgGet, (28, 28))
            # cv2.imshow("imgB", imgB)
            # cv2.waitKey(0)
            
            Temptuple = (imgGet, xx, yy, w, h, ss, area)  # 将图像及其坐标放在一个元组里面，然后再放进一个列表里面就可以访问了
            Borderlist.append(Temptuple)

    return Borderlist

def getImgMaxIdx(directory):
    maxIdx = 0
    for filename in os.listdir(directory):
        path = os.path.join(directory, filename)
        if os.path.isfile(path):
            items = str.split(filename, '.')
            if items[-1] == "png":
                idx = int(items[-2])
                if idx > maxIdx:
                    maxIdx = idx
    return maxIdx

class AccuratePosi:
    def __init__(self, d, x, y):
        self.d = int(d)
        self.x = int(x)
        self.y = int(y)

def getMinDeltaX(posList, thePos):
    res = []
    for pos in posList:
        deltaX = abs(pos.x - thePos.x)
        deltaY = abs(pos.y - thePos.y)
        if pos.x > thePos.x and deltaX <= 50 and deltaY <= 25:
            template = (pos, pos.x - thePos.x, pos.y - thePos.y)
            res.append(template)
    
    sorted(res, key=lambda p:p[1])
    return res

def dragToLeft(startPoi, distance):
    """鼠标往屏幕左边拖动一定距离, 视角往右移动
    """
    pg.moveTo(startPoi.x, startPoi.y, duration=1)
    pg.mouseDown(button='left')
    pg.dragRel(-distance, duration=1)
    pg.mouseUp(button='left')
    
# def dragToRight(distance):

class Model(torch.nn.Module):
    def __init__(self) :
        super(Model, self).__init__()
        self.conv1 = torch.nn.Sequential(torch.nn.Conv2d(1, 64, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.Conv2d(64, 128, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.MaxPool2d(2, 2)) 
        self.dense = torch.nn.Sequential(torch.nn.Linear(14*14*128, 1024),
                                         torch.nn.ReLU(),
                                         torch.nn.Dropout(p = 0.5),
                                         torch.nn.Linear(1024, 11))
        
    def forward(self, x) :
        x = self.conv1(x)
        x = x.view(-1, 14*14*128)
        x = self.dense(x)
        return x

def main(argv):
    ScreenSize = pg.size()
    print(ScreenSize)
    originalPoi = AccuratePosi(0xff, ScreenSize.width / 2 - (1600 / 2) + 4, ScreenSize.height / 2 - (900 / 2))
    screenImgRaw = pg.screenshot(region=[originalPoi.x, originalPoi.y, 1600, 900])
    screenImg = cv2.cvtColor(np.asarray(screenImgRaw), cv2.COLOR_RGB2BGR)
    cv2.imwrite("ScreenShot.jpg", screenImg)
    
    # imgFn = argv[0]
    # img = cv2.imread(imgFn)
    # imgShow = cv2.imread(imgFn)
    img = screenImg
    imgShow = cv2.cvtColor(np.asarray(screenImgRaw), cv2.COLOR_RGB2BGR)
    # cv2.rectangle(imgB, (x - 10, y - dd - 10), (x + w + 10, y + a + 10), (0, 255, 0), 2)
    cv2.rectangle(img, (0, 706), (565, 899), (0, 0, 0), cv2.FILLED) # 交流框
    cv2.rectangle(img, (0, 0), (1599, 157), (0, 0, 0), cv2.FILLED)  # 顶部框
    cv2.rectangle(img, (566, 808), (1599, 899), (0, 0, 0), cv2.FILLED) # 下部按钮
    cv2.rectangle(img, (1493, 157), (1599, 899), (0, 0, 0), cv2.FILLED) # 右侧栏
    cv2.rectangle(img, (1443, 741), (1599, 899), (0, 0, 0), cv2.FILLED) # 寻访按钮
    cv2.rectangle(img, (1443, 741), (1599, 899), (0, 0, 0), cv2.FILLED) # 寻访按钮
    cv2.rectangle(img, (13, 180), (80, 247), (0, 0, 0), cv2.FILLED) # 任务按钮
    # cv2.imshow("m", img)
    # if cv2.waitKey(0) == 113:
    #     exit()

    hsvImg = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    kernel = np.ones((3, 3), np.uint8)
    hsvImg1 =  cv2.GaussianBlur(hsvImg, (3, 3), 1)
    kernel = np.ones((1, 1), np.uint8)
    hsvImg2 = cv2.dilate(hsvImg1, kernel)
    kernel = np.ones((1, 1), np.uint8)
    hsvImg3 = cv2.erode(hsvImg2, kernel)
    
    h,s,v = cv2.split(hsvImg) # 分割后会得到HSV3个通道的值
    res = np.average(v) # 对V（明暗度）求平均值，也可以用np.sum/图片像素点总数来求平均值
    # print(f"v:{res}")
    # cv2.imshow("h", h)
    # cv2.imshow("s", s)
    # cv2.imshow("v", v)
    # cv2.imwrite("h.png", h)
    # cv2.imwrite("s.png", s)
    # cv2.imwrite("v.png", v)
    # cv2.waitKey(0)


    # 定义绿色的色调、饱和度和明度范围
    lower_green = np.array([0, 0, 150])
    upper_green = np.array([255, 45, 255])

    # 根据颜色范围创建掩码
    mask = cv2.inRange(hsvImg3, lower_green, upper_green)
    # 使用掩码提取绿色部分
    result = cv2.bitwise_and(img, img, mask=mask)

    # cv2.imshow('Original Image', img)
    # cv2.imshow('Mask', mask)
    # cv2.imshow('Result', result)
    # cv2.imwrite('result.png', result)
    
    grayImg = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)

    [ret,binImg] = cv2.threshold(grayImg, 165, 255, cv2.THRESH_BINARY)
    # cv2.imshow("binImg", binImg)
    # cv2.imwrite('binImg.png', binImg)
    # if cv2.waitKey(0) == 113:
    #     exit()

    blist = getContours(binImg, img)
    for (imgRes, x, y, w, h, s, a) in blist:
        locatePosStr = str("(%d,%d)" % (x, y))
        cv2.putText(img, str(locatePosStr),
                                (x,y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255))
        
    # cv2.imshow("contours", img)
    # if cv2.waitKey(0) == 113:
    #     exit()
    
    # 模型部分
    train_transform = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Grayscale(),
        transforms.Resize((28, 28)),
        transforms.ToTensor(),
    ])
    
    torch.serialization.add_safe_globals([Model])
    device = torch.device('cpu')
    model = Model().to(device)
    model.load_state_dict(torch.load('parameter.pkl', weights_only=False)) #load
    index_to_class = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'f']
    
    
    print("start test...")
    start_time = time.time()
    test_correct = 0
    
    # 最后一个系
    digit_count = [0,0,0,0,0,0,0,0,0,0,0]
    for i in range(0, 10):
        os.makedirs("data/SM/train4/" + str(i), exist_ok=True)
        
    os.makedirs("data/SM/train4/other", exist_ok=True)
    
    # 从各文件夹里获取最新图像的序号
    for i in range(0, 10):
        digit_count[i] = getImgMaxIdx("data/SM/train4/" + str(i))
    ## other文件夹也要获取
    digit_count[10] = getImgMaxIdx("data/SM/train4/other")
    
    # 创建准确数字集合
    AccurateDigitPosList = []
    AccurateDigitCount = 0
    blistLen = len(blist)
    blistCount = 0
    if len(blist) != 0:
        for (imgRes, x, y, w, h, s, a) in blist:
            blistCount+=1
            imgTmp = train_transform(imgRes)
            imgTmp = torch.unsqueeze(imgTmp, dim=0)
            with torch.no_grad():
                pre = model(imgTmp)
                output = torch.squeeze(pre)
                predict = torch.softmax(output, dim=0)
                predict_cla = torch.argmax(predict).numpy()
                print(predict_cla)
                print(index_to_class[predict_cla], predict[predict_cla].numpy())
                if predict[predict_cla].numpy() >= 0.90:
                    if index_to_class[predict_cla] == 'f':
                        print(str(index_to_class[predict_cla]) + " " + 
                            str("%.04f" % predict[predict_cla].numpy()) + " " +
                            str("%.02f" % a) ,
                            (x,y), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0))
                    else: # elif x >700 and x < 800 and y > 562 and y < 622:
                        print(str(x) + ',' + str(y) + ';' + str(index_to_class[predict_cla]) + " " + 
                            str("%.04f" % predict[predict_cla].numpy()) + " " +
                            str("%.02f" % a) , (x,y), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0))
                        cv2.rectangle(imgShow, (x, y), (x + s, y + s), color=(0,255,0), thickness=1)
                        cv2.putText(imgShow, str(index_to_class[predict_cla]) + " " + str("%.04f" % predict[predict_cla].numpy()),
                                    (x,y), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,0,0))
                        pos = AccuratePosi(int(index_to_class[predict_cla]), x, y)
                        AccurateDigitPosList.append(pos)
                        AccurateDigitCount += 1
                        
    
    ####### 目标:获取有效两位数的坐标; 输出:DoubleDigitPlotList #######
    DoubleDigitPlotList = []
    sortX = sorted(AccurateDigitPosList, key=lambda t:t.x)
    sortY = sorted(AccurateDigitPosList, key=lambda t:t.y)
    for i in range(len(sortY)):
        j = i + 1
        if (j >= len(sortY)):
            break
        posA = sortY[i]
        posB = sortY[j]
        deltaY = posB.y - posA.y
        if deltaY <= 20:
            if posA.x > posB.x:
                tmpPos = posA
                posA = posB
                posB = tmpPos
            
            if posA.d == 1:
                print("posA:(%d, %d), posB:(%d, %d)" %(posA.x, posA.y, posB.x, posB.y))
                xlist = getMinDeltaX(sortX, posA)
                print(len(xlist))
                if len(xlist) != 0:
                    cv2.circle(imgShow,  (posA.x + 23, posA.y), 2, (0,0,255), 2)
                    # cv2.rectangle(imgShow, (posA.x, posA.y), (posA.x + 23, posA.y + 23), color=(255,0,0), thickness=1)
                    # cv2.putText(imgShow, str(posA.x) + ',' + str(posA.y), (posA.x, posA.y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255))
                    
                
                for (pos, dx, dy) in xlist:
                    if dx != 0 and dx <= 25:
                        cv2.rectangle(imgShow, (pos.x, pos.y), (pos.x + 23, pos.y + 23), color=(0,0,255), thickness=2)
                        firstDigitIdx = AccurateDigitPosList.index(posA)
                        secondDigitIdx = AccurateDigitPosList.index(pos)
                        idxTemplate = (firstDigitIdx, secondDigitIdx)
                        DoubleDigitPlotList.append(idxTemplate)
                        break
    
    ## pgautogui初始化    
    ScreenSize = pg.size()
    print(ScreenSize)
    pg.moveTo(ScreenSize.width / 2, ScreenSize.height / 2)
    
    originalPoi = AccuratePosi(0xff, ScreenSize.width / 2 - (1600 / 2), ScreenSize.height / 2 - (900 / 2))
    tmpPoi = AccuratePosi(0xff, originalPoi.x + 1399, originalPoi.y + (900 / 2))
    # for i in range(10):
    #     dragToLeft(tmpPoi, 1000)
    
    ####### 目标:间隔1s点击指定位置; #######
    for (fidx, sidx) in DoubleDigitPlotList:
        firstDigit = AccurateDigitPosList[fidx]
        secondDigit = AccurateDigitPosList[sidx]
        pg.moveTo(originalPoi.x + secondDigit.x + 23, 
                  originalPoi.y + secondDigit.y, 
                  duration=1)
        pg.click(originalPoi.x + secondDigit.x + 23, 
                 originalPoi.y + secondDigit.y,
                 button='left')
       
    
    print("done")
    
    cv2.imshow("imgShow", imgShow)
    cv2.waitKey(0)

if __name__ == "__main__":
    main(sys.argv[1:])